import numpy as np
import pandas as pd
import statsmodels.api as sm
from sklearn.cluster import KMeans
from tqdm import tqdm

# ========= Step 1: 读取数据 =========
file_path = "预处理3——孕周数（小数）.xlsx"
df_male = pd.read_excel(file_path, sheet_name="男胎检测数据")

# 提取关键变量
df = df_male[['孕妇BMI', '检测孕周', 'Y染色体浓度']].dropna().copy()

# ========= Step 2: 定义蒙特卡洛模拟函数 =========
def monte_carlo_simulation(df, n_iter=1000, sigma=0.1, n_clusters=3):
    """
    蒙特卡洛模拟:
    - 对 Y 染色体浓度做 logit-normal 扰动
    - 重建逻辑回归模型 + BMI 分组
    - 统计每组孕周中位数
    """
    all_nipt_times = []

    for _ in tqdm(range(n_iter), desc="Running simulations"):
        df_sim = df.copy()

        # 避免 0/1，先做伪计数
        eps = 1e-6
        p = np.clip(df_sim['Y染色体浓度'].values, eps, 1 - eps)

        # logit 变换 + 正态扰动
        logit_p = np.log(p / (1 - p))
        logit_p_sim = logit_p + np.random.normal(0, sigma, size=len(p))
        p_sim = 1 / (1 + np.exp(-logit_p_sim))
        df_sim['Y扰动'] = p_sim

        # 达标标签
        df_sim['达标'] = (df_sim['Y扰动'] >= 0.04).astype(int)

        # 逻辑回归（BMI + 孕周）
        X_sim = sm.add_constant(df_sim[['孕妇BMI', '检测孕周']])
        y_sim = df_sim['达标']

        try:
            sm.Logit(y_sim, X_sim).fit(disp=False)
        except:
            continue  # 拟合失败就跳过

        # KMeans 分组
        kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=None)
        df_sim['BMI分组'] = kmeans.fit_predict(df_sim[['孕妇BMI']])

        # 每组孕周中位数
        g = df_sim.groupby('BMI分组').agg({
            '孕妇BMI': ['min', 'max', 'median'],
            '检测孕周': 'median'
        }).reset_index()

        if len(g) == n_clusters:
            all_nipt_times.append(list(g['检测孕周']['median']))

    return np.array([np.sort(x) for x in all_nipt_times if len(x) == n_clusters])


# ========= Step 3: 跑 1000 次模拟 =========
nipt_array = monte_carlo_simulation(df, n_iter=1000, sigma=0.1, n_clusters=3)

# ========= Step 4: 统计均值 + 95%CI =========
nipt_mean = np.nanmean(nipt_array, axis=0)
nipt_ci_lower = np.nanpercentile(nipt_array, 2.5, axis=0)
nipt_ci_upper = np.nanpercentile(nipt_array, 97.5, axis=0)

nipt_results = pd.DataFrame({
    'BMI分组': [1, 2, 3],
    '扰动后孕周均值': nipt_mean,
    '孕周95%CI下限': nipt_ci_lower,
    '孕周95%CI上限': nipt_ci_upper
})

# 保存结果
nipt_results.to_excel("NIPT蒙特卡洛模拟结果.xlsx", index=False)
print("✅ 模拟完成，结果已保存到 NIPT蒙特卡洛模拟结果.xlsx")
print(nipt_results)
